from cam.cam_base import CamBase
from env.env_base import EnvBase
import torch
import argparse
import json
from math import cos, sin
import pathlib
import time
import cv2
import numpy as np


class AffineCameraPhoto(EnvBase):
    def load_camera(
        self,
        f_x: float,
        f_y: float,
        u_0: float,
        v_0: float,
        dx: float,
        dy: float,
        h: float,
        w: float,
        **kwargs,
    ) -> None:
        """
        Arguments:
         - `f_x` : x轴焦距，单位是毫米
         - `f_y` : y轴焦距，单位是毫米
         - `u_0` : x轴方向上主点与原点的偏移，单位是像素
         - `v_0` : y轴方向上主点与原点的偏移，单位是像素
         - `dx` : x轴上单位像素对应的毫米数
         - `dy` : y轴上单位像素对应的毫米数
         - `h` : 相机在x轴方向上的分辨率
         - `w` : 相机在y轴方向上的分辨率
        """
        self._cam = CamBase(f_x, f_y, u_0, v_0, dx, dy, h, w, **kwargs)

    def photo_by_camera(self, path: str, scale: float, **kwargs) -> np.ndarray:
        imag_flat, p_s = self._load_2d(path=path, scale=scale, **kwargs)
        return self.cam._photo(p_s, imag_flat)


def args_parse():
    args_parser = argparse.ArgumentParser()
    args_parser.add_argument("--rx", type=float, default=0.01)
    args_parser.add_argument("--ry", type=float, default=0.01)
    args_parser.add_argument("--rz", type=float, default=0.01)
    args_parser.add_argument("--x", type=float, default=5)
    args_parser.add_argument("--y", type=float, default=5)
    args_parser.add_argument("--z", type=float, default=1)
    args_parser.add_argument("--scale", type=float, default=1)
    args_parser.add_argument("--mate_path", type=str, default="../check_board.png")
    args_parser.add_argument(
        "--cam_ind", type=int, default=1, help="corespond to the camera.jsonl indexes."
    )

    args = args_parser.parse_args()
    return args.__dict__


def R_parse(theta_x: float, theta_y: float, theta_z: float):
    rad = 3.1415926 / 180
    theta_x *= rad
    theta_y *= rad
    theta_z *= rad
    return (
        torch.tensor(
            [
                [1, 0, 0],
                [0, cos(theta_x), sin(theta_x)],
                [0, -sin(theta_x), cos(theta_x)],
            ],
            dtype=torch.float32,
        )
        @ torch.tensor(
            [
                [cos(theta_y), 0, -sin(theta_y)],
                [0, 1, 0],
                [sin(theta_y), 0, cos(theta_y)],
            ],
            dtype=torch.float32,
        )
        @ torch.tensor(
            [
                [cos(theta_z), sin(theta_z), 0],
                [-sin(theta_z), cos(theta_z), 0],
                [0, 0, 1],
            ],
            dtype=torch.float32,
        )
    )


def t_parse(x: float, y: float, z: float):
    return torch.tensor([x, y, z], dtype=torch.float32)


def p_parse(x: float, y: float, z: float, w: float):
    return torch.tensor([x, y, z, w], dtype=torch.float32)


def json_parse(cam_path: str, cam_ind: int) -> dict:
    if pathlib.Path(cam_path).exists():
        with open(cam_path, "r") as f:
            for line in f:
                cam = json.loads(line)
                if cam["index"] == cam_ind:
                    print(
                        "-----Using this chosen camera <- "
                        + cam["name"]
                        + f" -> :{cam_ind}-----"
                    )
                    return cam["meta"]
    else:
        raise FileNotFoundError(f"{cam_path} not found")


if __name__ == "__main__":
    args = args_parse()
    R = R_parse(args["rx"], args["ry"], args["rz"])
    t = t_parse(args["x"], args["y"], args["z"])
    E = AffineCameraPhoto(R, t)
    imag_flat, p_s = E._load_2d(args["mate_path"], args["scale"])

    cam = json_parse("./cam/cameras.jsonl", args["cam_ind"])
    E.load_camera(**cam)
    E.show()
    image = E.photo_by_camera(path=args["mate_path"], scale=args["scale"])

    if not pathlib.Path("./output").exists():
        pathlib.Path("./output").mkdir()
    name = time.strftime("%Y-%m-%d-%H-%M-%S")
    cv2.imwrite("./output/" + name + ".png", img=image)
    E.save("./output/" + name + ".pt")
